from abc import ABC, abstractmethod
from typing import Dict, Any, Optional

# 条件导入 FEniCSx 类型和 SimulationParameters
try:
    import dolfinx
    from dolfinx import fem, mesh # type: ignore
    from mpi4py import MPI # type: ignore
    FENICSX_AVAILABLE = True
    # 导入 SimulationParameters，假设它在项目的某个可访问路径
    # 例如: from ..core.simulation_parameters import SimulationParameters
    # 为简单起见，此处暂时使用占位符，实际项目中应正确导入
    class SimulationParameters: # 占位符
        pass
    class FEMEntities: # 占位符，代表 fem_entities 的结构
        pass

except ImportError:
    FENICSX_AVAILABLE = False
    # 如果 FEniCSx 不可用，则定义用于类型提示的虚拟类型
    class fem:
        class Function:
            pass
    class mesh:
        class Mesh:
            pass
    class MPI: # type: ignore
        class Comm:
            pass
    class SimulationParameters: # 占位符
        pass
    class FEMEntities: # 占位符
        pass


class OutputManagerBase(ABC):
    """
    输出和后处理管理器的抽象基类。

    该类为在模拟的不同阶段执行输出（例如，保存结果到文件）
    或后处理操作（例如，计算统计数据、可视化）提供了一个通用接口。
    """

    def __init__(self, 
                 comm: Optional[MPI.Comm], 
                 params: SimulationParameters, 
                 fem_entities: Dict[str, Any], 
                 output_dir: str):
        """
        初始化 OutputManagerBase。

        参数:
            comm (Optional[MPI.Comm]): MPI 通信器。
            params (SimulationParameters): 包含所有模拟参数的 Pydantic 模型实例。
            fem_entities (Dict[str, Any]): 包含基本 FEniCSx FEM 实体的字典，
                                           例如网格、函数空间和解函数。
            output_dir (str): 输出文件存放的目录。
        """
        self.comm = comm
        self.params = params
        self.fem_entities = fem_entities
        self.output_dir = output_dir
        
        if FENICSX_AVAILABLE and self.comm is None:
            # 如果 FEniCSx 可用但未提供 comm，尝试获取默认通信器
            # 这只是一个示例，具体行为可能需要根据项目调整
            try:
                from mpi4py import MPI as mpi_real
                self.comm = mpi_real.COMM_WORLD
                if self.comm.rank == 0:
                    print("信息: OutputManagerBase 未提供 MPI comm，已自动设为 COMM_WORLD。")
            except ImportError:
                if self.comm is None and self.comm.rank == 0: # type: ignore
                     print("警告: mpi4py 未找到，MPI comm 为 None。某些输出功能可能受限。")


        print(f"信息: {self.__class__.__name__} 初始化完成。输出目录: '{self.output_dir}'")

    @abstractmethod
    def is_output_step(self, iteration_count: int, current_time: float) -> bool:
        """
        判断当前迭代/时间是否需要执行输出或后处理。

        参数:
            iteration_count (int): 当前的迭代次数。
            current_time (float): 当前的模拟时间。

        返回:
            bool: 如果需要执行输出/后处理，则为 True，否则为 False。
        """
        pass

    @abstractmethod
    def execute_step(self, iteration_count: int, current_time: float) -> None:
        """
        执行当前步骤的输出和/或后处理操作。

        仅当 `is_output_step` 返回 True 时，此方法才应被调用。

        参数:
            iteration_count (int): 当前的迭代次数。
            current_time (float): 当前的模拟时间。
        """
        pass

    @abstractmethod
    def finalize(self) -> None:
        """
        执行所有最终的输出/后处理操作。

        例如，关闭文件句柄，完成最终的统计计算等。
        此方法通常在模拟结束时调用一次。
        """
        pass

if __name__ == '__main__':
    # 这是一个演示如何使用该基类的示例
    
    # 模拟 SimulationParameters, MPI.Comm, 和 fem_entities
    class MockMPIComm: # 模拟 MPI.Comm
        rank = 0
        size = 1
        def bcast(self, data, root=0): return data
        def Barrier(self): pass
        def allreduce(self, data, op=None): return data # 简化 allreduce

    class MockSimParams(SimulationParameters):
        def __init__(self):
            self.output_frequency = 2 # MyVTKOutput 的输出频率
            self.stats_calc_frequency = 1 # SolutionStatisticsCalculator 的频率
            self.csv_log_frequency = 1    # CSVDataLogger 的频率
            self.custom_data_point = 100  # CSVDataLogger 回调要用的示例参数
            # 其他参数...

    # ----- 为 CSVDataLogger 定义一个回调函数 -----
    def example_csv_data_callback(iter_count, current_time, fem_entities, params_obj):
        """示例回调函数，用于提取要记录到 CSV 的数据。"""
        # 从 fem_entities 获取一些模拟数据 (例如，一个标量函数在某一点的值)
        # 从 params_obj 获取一些参数
        # 这里仅为演示，返回一些构造的数据
        val_from_fem = None
        if FENICSX_AVAILABLE and "solution_phi_s" in fem_entities:
            # 假设我们想获取函数在某个特定点的值 (这里简化，用均值代替)
            # 实际应用中可能需要更复杂的插值或点评估
            try:
                # local_values = fem_entities["solution_phi_s"].x.array
                # val_from_fem = np.mean(local_values) if local_values.size > 0 else 0.0
                # 更简单的方式：直接获取一个模拟值，或者如果可以，用一个全局统计量
                if hasattr(fem_entities["solution_phi_s"], 'x'):
                     val_from_fem = fem_entities["solution_phi_s"].x.array[0] # 取第一个节点值作为示例
            except Exception:
                val_from_fem = -1.0 # 出错时的默认值

        param_val = params_obj.custom_data_point if hasattr(params_obj, 'custom_data_point') else -1
        return [iter_count, f"{current_time:.3f}", val_from_fem, param_val]
    # ----- 回调函数定义结束 -----

    mock_params_instance = MockSimParams()
    comm_instance = MockMPIComm() # 默认使用模拟的 Comm
    
    # 根据 FENICSX_AVAILABLE 准备 fem_entities
    if FENICSX_AVAILABLE:
        try:
            from mpi4py import MPI as mpi_real # 尝试导入真实的 MPI
            comm_instance = mpi_real.COMM_WORLD
            print(f"信息: 使用真实的 MPI.COMM_WORLD (rank {comm_instance.rank}/{comm_instance.size})")

            import numpy as np # 仅当 FEniCSx 可用时导入
            domain = dolfinx.mesh.create_unit_square(comm_instance, 5, 5) # type: ignore
            V_phi = dolfinx.fem.functionspace(domain, ("Lagrange", 1)) # type: ignore
            V_temp = dolfinx.fem.functionspace(domain, ("Lagrange", 1)) # type: ignore
            
            solution_phi_s_func = dolfinx.fem.Function(V_phi, name="phi_s_example") # type: ignore
            solution_phi_s_func.interpolate(lambda x: x[0] + x[1] + np.random.rand(x.shape[1])) # type: ignore
            
            solution_temp_func = dolfinx.fem.Function(V_temp, name="temperature_example") # type: ignore
            solution_temp_func.interpolate(lambda x: 300 + 10 * x[0] - 5 * x[1] + 5 * np.random.rand(x.shape[1])) # type: ignore
            
            mock_fem_entities_instance = {
                "solution_phi_s": solution_phi_s_func,
                "solution_temp": solution_temp_func,
                "mesh": domain
            }
            print("信息: 已创建基于 FEniCSx 的模拟 FEM 实体。")
        except Exception as e:
            if comm_instance.rank == 0:
                print(f"警告: 创建 FEniCSx 实体时出错: {e}。将回退到非 FEniCSx 模拟实体。")
            FENICSX_AVAILABLE = False # 如果真实 FEniCSx 设置失败，则标记为不可用

    if not FENICSX_AVAILABLE: # 如果 FEniCSx 不可用或上面初始化失败
        comm_instance = MockMPIComm() # 确保使用模拟 Comm
        mock_fem_entities_instance = { # type: ignore
            "solution_phi_s": "这是一个模拟的phi_s解函数",
            "solution_temp": "这是一个模拟的温度解函数",
            "mesh": "这是一个模拟的网格对象"
        }
        if comm_instance.rank == 0:
            print("信息: 使用非 FEniCSx 的模拟 FEM 实体。")


    # --- 子类演示 ---
    # 导入子类 (确保它们在 Python 路径中，或者调整导入)
    # 假设这些处理器在 output_handlers 子目录中
    try:
        from ..output_handlers.vtk_output_handler import VTKOutputHandler
        from ..output_handlers.solution_statistics_calculator import SolutionStatisticsCalculator
        from ..output_handlers.csv_data_logger import CSVDataLogger
        from ..output_handlers.pyvista_visualizer import PyVistaVisualizer # 如果要测试
        HANDLERS_IMPORTED = True
    except ImportError:
        if comm_instance.rank == 0:
            print("\n警告: 无法从 output_handlers 导入处理器类。确保它们在正确的路径并且 __init__.py 存在。")
            print("将跳过 OutputManagerBase 的子类演示。")
        HANDLERS_IMPORTED = False

    if HANDLERS_IMPORTED:
        output_directory = "results_output_manager_base_demo"
        if comm_instance.rank == 0 and FENICSX_AVAILABLE:
             print(f"\n所有输出将尝试写入 '{output_directory}' (如果适用)。")
        
        # 1. VTKOutputHandler 演示
        if comm_instance.rank == 0: print("\n--- VTKOutputHandler 演示 ---")
        vtk_handler = VTKOutputHandler(
            comm=comm_instance, 
            params=mock_params_instance, 
            fem_entities=mock_fem_entities_instance, 
            output_dir=output_directory,
            filename_prefix="demo_solution",
            output_frequency=mock_params_instance.output_frequency
        )

        # 2. SolutionStatisticsCalculator 演示
        if comm_instance.rank == 0: print("\n--- SolutionStatisticsCalculator 演示 ---")
        stats_calculator = SolutionStatisticsCalculator(
            comm=comm_instance,
            params=mock_params_instance,
            fem_entities=mock_fem_entities_instance,
            output_dir=output_directory,
            stats_frequency=mock_params_instance.stats_calc_frequency,
            functions_to_analyze=["solution_phi_s", "solution_temp", "non_existent_func"] # 测试存在的和不存在的
        )

        # 3. CSVDataLogger 演示
        if comm_instance.rank == 0: print("\n--- CSVDataLogger 演示 ---")
        csv_logger = CSVDataLogger(
            comm=comm_instance,
            params=mock_params_instance,
            fem_entities=mock_fem_entities_instance,
            output_dir=output_directory,
            csv_filename="simulation_log.csv",
            log_frequency=mock_params_instance.csv_log_frequency,
            headers=["Iteration", "Time (s)", "Phi_s_Value_Example", "CustomParam"],
            data_extraction_callback=example_csv_data_callback
        )
        
        # (可选) PyVistaVisualizer 演示
        # pyvista_handler = PyVistaVisualizer(
        #     comm=comm_instance,
        #     params=mock_params_instance,
        #     fem_entities=mock_fem_entities_instance,
        #     output_dir=output_directory,
        #     visualize_frequency=3, # 每 3 次迭代可视化一次
        #     interactive_visualization=False # 设为 True 以尝试交互式窗口 (可能不工作)
        # )

        # 将所有处理器放入列表
        output_processors = [vtk_handler, stats_calculator, csv_logger] #, pyvista_handler]

        # 模拟主循环
        if comm_instance.rank == 0: print("\n--- 开始模拟循环演示 ---")
        num_iterations_demo = 5
        for i in range(num_iterations_demo):
            current_t_demo = float(i) * 0.1 # 模拟时间
            if comm_instance.rank == 0:
                print(f"\n模拟迭代 {i}, 时间 {current_t_demo:.2f}s")

            # 在每次迭代中更新模拟解函数的值 (如果 FEniCSx 可用)
            if FENICSX_AVAILABLE and "solution_phi_s" in mock_fem_entities_instance:
                mock_fem_entities_instance["solution_phi_s"].interpolate( # type: ignore
                    lambda x: x[0] + x[1] + float(i) + np.random.rand(x.shape[1]) # type: ignore
                )
                mock_fem_entities_instance["solution_temp"].interpolate( # type: ignore
                    lambda x: 300 + 10*x[0] - 5*x[1] + float(i) + 5*np.random.rand(x.shape[1]) # type: ignore
                )


            for processor in output_processors:
                if processor.is_output_step(i, current_t_demo):
                    processor.execute_step(i, current_t_demo)
            
            if FENICSX_AVAILABLE: comm_instance.Barrier() # 同步 MPI 进程 (如果使用真实 MPI)

        # 结束所有处理器
        if comm_instance.rank == 0: print("\n--- 模拟循环结束, 执行 finalize ---")
        for processor in output_processors:
            processor.finalize()
            if FENICSX_AVAILABLE: comm_instance.Barrier()
    
    if comm_instance.rank == 0:
        print("\nOutputManagerBase 及其子类演示完成。")
        if FENICSX_AVAILABLE and HANDLERS_IMPORTED:
            print(f"请检查目录 '{output_directory}' 中的输出文件 (VTK, CSV, JSON)。") 